import pandas as pd
import xlrd
import xlwt
import datetime
import tushare as ts
ts.set_token('883c1e2da92c0f68194fe2c4b4aa0e4da880b41ef62ff556f2da00e7')
pro = ts.pro_api()
def charge_time(x):
    df = pro.trade_cal(exchange='', start_date=x, end_date=x)
    if df['is_open'].values > 0:
        return x
    else:
        date1 =datetime.datetime.strptime(x, '%Y%m%d')
        date2 = date1 + datetime.timedelta(days=1)
        date2 = date2.strftime('%Y%m%d')
        return charge_time(date2)

df = pd.read_hdf('stock_data.h5', key='data')
df = df.swaplevel('trade_date', 'ts_code').sort_index(level=0)

el = xlrd.open_workbook('peigu_data.xlsx')
sheet = el.sheet_by_name('Sheet2')

cols = sheet.col_values(0)
df1 = pd.DataFrame(cols)
list1 = []
list2 = []

for i, row in df1.iterrows():
    if i <= 2:
        continue
    code = row[0]

    date = xlrd.xldate_as_tuple(sheet.cell(i, 7).value, 0)
    value = datetime.datetime(*date)
    stock_date = value.strftime('%Y%m%d')

    date1 = charge_time(stock_date)

    df2 = df.ix[code]
    close = df2.loc[date1]['close']
    list1.append(close)
    df2.reset_index(inplace=True)
    df3 = df2[df2.trade_date >= date1]

    pct_chg = round((df3['close'].values[4] - df3['close'].values[0])/df3['close'].values[0],4)

    list2.append(pct_chg)



print(list1,list2)


wb = openpyxl.load_workbook('peigu_data.xlsx')
ws = wb.worksheets[0]

ws.insert_cols(3)
for index, row in enumerate(ws.rows):
    if index == 0:
        row[2].value = '股价'
    else:
        row[2].value = list[index]

wb.save('peigu_data1.xlsx')
